#include <cuda_runtime.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>
#include "prof.h"

// 错误处理宏
#define CUDA_CHECK(call) \
    do { \
        cudaError_t err = (call); \
        if (err != cudaSuccess) { \
            fprintf(stderr, "CUDA错误: %s 在 %s:%d\n", cudaGetErrorString(err), __FILE__, __LINE__); \
            exit(EXIT_FAILURE); \
        } \
    } while (0)

// 计时函数（毫秒）
double get_time_ms() {
    struct timespec ts;
    clock_gettime(CLOCK_MONOTONIC, &ts);
    return ts.tv_sec * 1000.0 + ts.tv_nsec / 1000000.0;
}

// 统计信息结构体
typedef struct {
    char name[32];      // 算法名称
    int runs;           // 运行次数
    double total_time;  // 总时间 (ms)
    double avg_time;    // 平均时间 (ms)
    double max_time;    // 最大时间 (ms)
    double min_time;    // 最小时间 (ms)
    double total_htod;  // 总htod时间 (ms)
    double total_dtoh;  // 总dtoh时间 (ms)
} Stats;

// CUDA核函数：一维FIR滤波
__global__ void fir_filter_kernel(float *d_input, float *d_filter, float *d_output, 
                                 int signal_len, int filter_len, int num_groups) {
    int group_idx = blockIdx.x;                // 每个block处理一组信号
    int tid = threadIdx.x;                     // 线程ID
    int signal_start = group_idx * signal_len; // 当前组信号的起始位置

    if (group_idx >= num_groups) return;

    // 每个线程计算一个输出点
    for (int i = tid; i < signal_len + filter_len - 1; i += blockDim.x) {
        float sum = 0.0f;
        for (int j = 0; j < filter_len; j++) {
            int input_idx = i - j + signal_start;
            if (input_idx >= signal_start && input_idx < signal_start + signal_len) {
                sum += d_input[input_idx] * d_filter[j];
            }
        }
        d_output[signal_start + i] = sum;
    }
}

int main(int argc, char** argv) {
    // 参数设置
    int num_groups = 10000;         // 信号组数
    int signal_len = 1024;        // 每组信号长度
    int filter_len = 256;         // 滤波器长度 (256阶)

    // 获取运行次数n
    if (argc != 2) {
        fprintf(stderr, "用法: %s <运行次数>\n", argv[0]);
        return 1;
    }
    int n = atoi(argv[1]);
    if (n <= 0) {
        fprintf(stderr, "运行次数必须为正整数。\n");
        return 1;
    }

    profInit();//初始化
    
    profSetSectName(1, "htod 时间");//根据需要设置计时段名称
    profSetSectName(2, "compute_time");
    profSetSectName(3, "dtoh 时间");

    int input_size = num_groups * signal_len;           // 输入数据总大小
    int output_size = num_groups * (signal_len + filter_len - 1); // 输出数据总大小

    // 分配主机内存
    float *h_input = (float *)malloc(input_size * sizeof(float));
    float *h_filter = (float *)malloc(filter_len * sizeof(float));
    float *h_output = (float *)malloc(output_size * sizeof(float));

    // 随机生成输入数据和滤波器系数
    srand(time(NULL));
    for (int i = 0; i < input_size; i++) {
        h_input[i] = (float)rand() / RAND_MAX;
    }
    for (int i = 0; i < filter_len; i++) {
        h_filter[i] = (float)rand() / RAND_MAX; // 简化为随机系数，可替换为低通滤波器设计
    }

    // 分配设备内存
    float *d_input, *d_filter, *d_output;
    CUDA_CHECK(cudaMalloc(&d_input, input_size * sizeof(float)));
    CUDA_CHECK(cudaMalloc(&d_filter, filter_len * sizeof(float)));
    CUDA_CHECK(cudaMalloc(&d_output, output_size * sizeof(float)));

    // 复制滤波器到设备（只需一次）
    CUDA_CHECK(cudaMemcpy(d_filter, h_filter, filter_len * sizeof(float), cudaMemcpyHostToDevice));

    // 设置CUDA网格和块
    int threads_per_block = 256;
    int blocks = num_groups;

    profEnterSect(0);
    // 运行n次
    for (int run = 0; run < n; run++) {

        // htod: 复制输入数据到设备
        profEnterSect(1);
        CUDA_CHECK(cudaMemcpy(d_input, h_input, input_size * sizeof(float), cudaMemcpyHostToDevice));
        profLeaveSect(1);

        // 计算
        profEnterSect(2);
        fir_filter_kernel<<<blocks, threads_per_block>>>(d_input, d_filter, d_output, 
                                                        signal_len, filter_len, num_groups);
        CUDA_CHECK(cudaDeviceSynchronize());
        profLeaveSect(2);

        // dtoh: 复制结果回主机
        profEnterSect(3);
        CUDA_CHECK(cudaMemcpy(h_output, d_output, output_size * sizeof(float), cudaMemcpyDeviceToHost));
        profLeaveSect(3);
    }
    profLeaveSect(0);
    profPrintInfo();

    // 释放内存
    free(h_input);
    free(h_filter);
    free(h_output);
    CUDA_CHECK(cudaFree(d_input));
    CUDA_CHECK(cudaFree(d_filter));
    CUDA_CHECK(cudaFree(d_output));

    return 0;
}